package tsj.ai.alibaba.data.analysis.node;

import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;
import tsj.ai.alibaba.data.analysis.enums.TextType;
import tsj.ai.alibaba.data.analysis.pojo.KeywordExtractionResult;
import tsj.ai.alibaba.data.analysis.prompt.PromptConstant;
import tsj.ai.alibaba.data.analysis.service.processing.QueryProcessingService;
import tsj.ai.alibaba.data.analysis.service.llm.BlockLlmService;
import tsj.ai.alibaba.data.analysis.util.ChatResponseUtil;
import tsj.ai.alibaba.data.analysis.util.StateUtil;
import tsj.ai.alibaba.data.analysis.util.FluxUtil;

import java.util.*;
import java.util.function.Consumer;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.stream.Collectors;

import static tsj.ai.alibaba.data.analysis.constant.Constant.*;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;

/**
 * @author taoshujian
 * @version 1.0
 * @since 2025-11-13 15:00:18
 */
@Slf4j
@Component
@AllArgsConstructor
public class KeywordExtractNode  implements NodeAction {

    @Autowired
    private BlockLlmService blockLlmService;

    private final QueryProcessingService queryProcessingService;


    @Override
    public Map<String, Object> apply(OverAllState state) throws Exception {
        String input = StateUtil.getStringValue(state, QUERY_REWRITE_NODE_OUTPUT,
                StateUtil.getStringValue(state, INPUT_KEY));

        try {
            Map<String, Object> resultMap = new HashMap<>();
            String question = expandQuestion(input);
            if(Objects.isNull(question) || question.isEmpty()){
                log.warn("问题拓展失败，使用原问题");
                question = input;
            }

            // 解析为List<String>
            ObjectMapper objectMapper = new ObjectMapper();
            List<String> questions = objectMapper.readValue(question,new TypeReference<List<String>>(){});
            //String result = processQuestions(questions, AGENTId);

            //问题变体处理
            List<KeywordExtractionResult> resultList = new ArrayList<>(questions.size());
            // 使用传统for循环替代forEach
            for (String q : questions) {
                List<String> keywords = queryProcessingService.extractKeywords(q);
                resultList.add(new KeywordExtractionResult(q, null, keywords));
            }

            // 合并关键词
            List<String> mergedKeywords = mergeKeywords(resultList, input);
            log.info("成功合并关键词: {}", mergedKeywords);

            return Map.of(KEYWORD_EXTRACT_NODE_OUTPUT, mergedKeywords);

        }
        catch (Exception e) {
            log.error("增强关键词提取失败{}", e.getMessage());
            throw new RuntimeException(e);
        }
    }


    /**
     * 问题拓展
     * @param query
     * @return
     */
    public String expandQuestion(String query) {
        log.info("问题拓展，原问题: {}", query);
        String responseFlux = null;
        try {
            Map<String, Object> params = new HashMap<>();
            params.put("question", query);

            String prompt = PromptConstant.getQuestionExpansionPromptTemplate().render(params);

            responseFlux = blockLlmService.chat(prompt);
            log.info("问题拓展结果: {}", responseFlux);
        }
        catch (Exception e) {
            log.warn("问题拓展失败: {}", e.getMessage());
            e.printStackTrace();
        }
        return responseFlux;
    }

    /**
     * 合并来自多个问题变体的证据，消除重复
     * @param extractionResults list of extraction results
     * @return merged evidence list
     */
    private List<String> mergeEvidences(List<KeywordExtractionResult> extractionResults) {
        Set<String> mergedEvidences = new HashSet<>();

        extractionResults.stream()
                .filter(KeywordExtractionResult::isSuccessful)
                .forEach(result -> mergedEvidences.addAll(result.getEvidences()));

        return new ArrayList<>(mergedEvidences);
    }


    /**
     * 合并多个问题变体中的关键字，消除重复并保持原始问题关键字的优先级
     * @param extractionResults 提取结果列表
     * @param originalQuestion 原始问题
     * @return 合并关键字列表
     */
    private List<String> mergeKeywords(List<KeywordExtractionResult> extractionResults, String originalQuestion) {
        if (extractionResults.isEmpty()) {
            return List.of();
        }

        Set<String> mergedKeywords = new LinkedHashSet<>();

        extractionResults.stream()
                .filter(result -> result.isSuccessful() && result.getQuestion().equals(originalQuestion))
                .findFirst()
                .ifPresent(result -> mergedKeywords.addAll(result.getKeywords()));

        extractionResults.stream()
                .filter(result -> result.isSuccessful() && !result.getQuestion().equals(originalQuestion))
                .forEach(result -> mergedKeywords.addAll(result.getKeywords()));

        return new ArrayList<>(mergedKeywords);
    }

    /**
     * 处理问题变体
     */
    private String processQuestions(List<String> questions, String agentId) {
        List<KeywordExtractionResult> resultList = new ArrayList<>(questions.size());


        // 使用传统for循环替代forEach
        for (int i = 0; i < questions.size(); i++) {
            String question = questions.get(i);
            List<String> evidences = queryProcessingService.extractKeywords(question);

            // 直接处理关键词提取，不使用Flux
            try {
                List<String> keywords = extractKeywordsDirectly(question, evidences);
                log.info("成功从问题变体提取关键词: 问题=\"{}\", 关键词={}", question, keywords);
                resultList.add(new KeywordExtractionResult(question, evidences, keywords));
            } catch (Exception e) {
                log.warn("从问题变体提取关键词失败: 问题={}", question, e);
                resultList.add(new KeywordExtractionResult(question, false));
            }
        }
        // 返回处理结果的字符串表示
        return formatResults(resultList);
    }

    // 新增方法：直接提取关键词（同步方式）
    private List<String> extractKeywordsDirectly(String question, List<String> evidences) {
        // 这里需要实现同步的关键词提取逻辑
        // 根据实际业务需求实现
        return new ArrayList<>(); // 占位符，需替换为实际实现
    }

    // 新增方法：格式化结果为字符串
    private String formatResults(List<KeywordExtractionResult> results) {
        StringBuilder sb = new StringBuilder();
        for (KeywordExtractionResult result : results) {
            sb.append(result.toString()).append(";");
        }
        return sb.toString();
    }



}
